from scipy import ndimage as ndi
import matplotlib.pyplot as plt

from skimage.segmentation import watershed
from skimage.morphology import disk
from skimage import data
from skimage.filters import rank
from skimage.color import rgb2gray
from skimage.util import img_as_ubyte

# Load image and convert to grayscale
img = data.astronaut()
img_gray = rgb2gray(img)
image = img_as_ubyte(img_gray)

# Calculate the local gradients and create markers
markers = rank.gradient(image, disk(5)) < 20
markers = ndi.label(markers)[0]
gradient = rank.gradient(image, disk(2))

# Apply Watershed algorithm
labels = watershed(gradient, markers)

# Plotting
fig, axes = plt.subplots(nrows=2, ncols=2, figsize=(8, 8), sharex=True, sharey=True, subplot_kw={'box_aspect':1})
ax = axes.ravel()

ax[0].imshow(image, cmap='gray', interpolation='nearest')
ax[0].set_title("Original")

ax[1].imshow(gradient, cmap='viridis', interpolation='nearest')
ax[1].set_title("Local Gradient")

ax[2].imshow(markers, cmap='viridis', interpolation='nearest')
ax[2].set_title("Markers")

ax[3].imshow(image, cmap='gray', interpolation='nearest')
ax[3].imshow(labels, cmap='viridis', interpolation='nearest', alpha=0.7)
ax[3].set_title("Segmented")

for a in ax:
    a.axis('off')

fig.tight_layout()
plt.show()
